package com.leetcode.根据数据结构分类.二叉树;

import com.leetcode.datastructure.TreeNode;

/**
 * @author: ZhouBert
 * @date: 2021/2/25
 * @description: 二叉查找树中第 K 小的元素 II
 * https://leetcode-cn.com/leetbook/read/high-frequency-algorithm-exercise/5hhxs5/
 * 给定一个二叉搜索树，我们希望找到其中第 k 小的元素。
 * <p>
 * 在这道题目中，除了原始的二叉搜索树 root 以外，你还会得到一个和其结构完全一致的二叉树 nodenum_root，
 * 树上节点的值代表以该节点为根的子树的节点数量。
 */
public class LC二叉搜索树中第K小的元素II {

	public static void main(String[] args) {
		LC二叉搜索树中第K小的元素II action = new LC二叉搜索树中第K小的元素II();
		test1(action);
		test2(action);
		test3(action);
	}

	public static void test1(LC二叉搜索树中第K小的元素II action) {
		// 3
		TreeNode root = TreeNode.stringToTreeNode("[5,3,8,2,3]");
		TreeNode nodenum_root = TreeNode.stringToTreeNode("[5,3,1,1,1]");
		int k = 5;
		int res = action.kthSmallestII(root, nodenum_root, k);
		System.out.println("res = " + res);
	}

	public static void test2(LC二叉搜索树中第K小的元素II action) {
		// 14
		TreeNode root = TreeNode.stringToTreeNode("[12,5,22,4,11,15,25,2,5,9,12,14,17,24,29,2,3,4,null,6,10,null,null,13,15,15,21,24,25,29,30,2,null,3,null,null,null,6,8,null,10,12,14,null,15,15,17,19,21,22,null,null,null,28,null,null,30,1,2,3,3,null,6,7,9,null,11,12,12,null,14,null,null,null,null,16,17,19,20,null,null,22,23,26,null,30,null,null,null,null,2,null,3,null,null,null,null,null,8,8,null,null,null,null,null,null,null,null,null,null,null,null,null,19,null,19,21,null,null,null,null,null,27,null,null,null,null,null,null,null,null,null,9,18,19,null,null,null,null,null,27,null,null,18,19]");
		TreeNode nodenum_root = TreeNode.stringToTreeNode("[76,30,45,14,15,28,16,11,2,13,1,9,18,6,9,5,5,1,null,9,3,null,null,6,2,5,12,4,1,5,3,4,null,4,null,null,null,2,6,null,2,3,2,null,1,1,3,10,1,3,null,null,null,4,null,null,2,1,2,2,1,null,1,2,3,null,1,1,1,null,1,null,null,null,null,1,1,6,3,null,null,1,1,3,null,1,null,null,null,null,1,null,1,null,null,null,null,null,1,2,null,null,null,null,null,null,null,null,null,null,null,null,null,5,null,1,1,null,null,null,null,null,2,null,null,null,null,null,null,null,null,null,1,3,1,null,null,null,null,null,1,null,null,1,1]");
		int k = 36;
		int res = action.kthSmallestII(root, nodenum_root, k);
		System.out.println("res = " + res);
	}

	public static void test3(LC二叉搜索树中第K小的元素II action) {
		// 3
		TreeNode root = TreeNode.stringToTreeNode("[1,null,15,8,21,3,10,17,26,1,7,null,11,16,19,22,30,null,1,3,null,null,13,null,17,19,null,21,22,26,30,null,null,null,5,null,null,17,null,17,null,null,21,null,23,null,null,null,null,5,7,null,null,null,null,null,22,23,25]");
		TreeNode nodenum_root = TreeNode.stringToTreeNode("[34,null,33,12,20,8,3,7,12,2,5,null,2,3,3,8,3,null,1,4,null,null,1,null,2,2,null,3,4,1,1,null,null,null,3,null,null,1,null,1,null,null,2,null,3,null,null,null,null,1,1,null,null,null,null,null,1,1,1]");
		int k = 5;
		int res = action.kthSmallestII(root, nodenum_root, k);
		System.out.println("res = " + res);
	}


	/**
	 * 如果没有 nodenum_root 会考虑使用中序遍历 暂存到数组 来解决，这样是 O(n) 的；
	 * 有了 nodenum_root 就考虑利用节点的个数，达到 O(h) 的复杂度
	 * --
	 * 二叉树的题目还是很适用递归的
	 * 由于题目木有对于找不到时的规定，暂时认为必定能够找到
	 * --
	 * 需要注意：当需要去右子树寻找时，得修改 k 进行下一次递归。
	 * @param root
	 * @param nodenum_root
	 * @param k
	 * @return
	 */
	public int kthSmallestII(TreeNode root, TreeNode nodenum_root, int k) {
		if (nodenum_root.left == null){
			if (k == 1){
				return root.val;
			}
			return kthSmallestII(root.right, nodenum_root.right, k - 1);
		}
		if (nodenum_root.left.val == k) {
			//说明 答案就在 root.left.right.right... 中
			TreeNode cur = root.left;
			while (cur.right != null) {
				cur = cur.right;
			}
			return cur.val;
		} else if (nodenum_root.left.val == k - 1) {
			//说明就是根节点
			return root.val;
		} else if (nodenum_root.left.val < k) {
			//说明答案在 right 子树中
			return kthSmallestII(root.right, nodenum_root.right, k - nodenum_root.left.val - 1);
		} else {
			return kthSmallestII(root.left, nodenum_root.left, k);
		}

	}
}
